package ch.ethz.fcl.process.gradient;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;

import javax.imageio.ImageIO;

import ch.ethz.fcl.tools.ImageTools;

/**
 * reference:
 * http://www.calvin.edu/academic/engineering/2011-12-team1/Code/Java/
 * CannyEdgeDetector.java
 * 
 * @author crazyfrog
 * 
 */
public class Gradient {
	private BufferedImage image;

	private final static float GAUSSIAN_CUT_OFF = 0.005f;

	private int height;
	private int width;

	private float[] xConv;
	private float[] yConv;
	private float[] xGradient;
	private float[] yGradient;

	private float gaussianKernelRadius;
	private int gaussianKernelWidth;

	private int kwidth;
	private float[] kernel;
	private float[] diffKernel;

	public Gradient() {
		initKernel();
	}

	public Gradient(BufferedImage image) {
		this.image = image;
		this.width = image.getWidth();
		this.height = image.getHeight();
		initKernel();
	}

	public void update(BufferedImage image) {
		this.image = image;
		height = image.getHeight();
		width = image.getWidth();
		process();
	}

	private void process() {
		initArrays();
		computeGradients(gaussianKernelRadius, gaussianKernelWidth);
	}

	private void initKernel() {
		gaussianKernelRadius = 1f;
		gaussianKernelWidth = 3;

		// generate the gaussian convolution masks
		kernel = new float[gaussianKernelWidth];
		diffKernel = new float[gaussianKernelWidth];

		for (kwidth = 0; kwidth < gaussianKernelWidth; kwidth++) {
			float g1 = gaussian(kwidth, gaussianKernelRadius);
			if (g1 <= GAUSSIAN_CUT_OFF && kwidth >= 2)
				break;
			float g2 = gaussian(kwidth - 0.5f, gaussianKernelRadius);
			float g3 = gaussian(kwidth + 0.5f, gaussianKernelRadius);
			kernel[kwidth] = (g1 + g2 + g3)
					/ 3f
					/ (2f * (float) Math.PI * gaussianKernelRadius * gaussianKernelRadius);
			diffKernel[kwidth] = g3 - g2;
		}

		// System.out.println("kernel width is " + kwidth);
	}

	// private utility methods
	private void initArrays() {
		xConv = new float[width * height];
		yConv = new float[width * height];
		xGradient = new float[width * height];
		yGradient = new float[width * height];
	}

	// NOTE: The elements of the method below (specifically the technique for
	// non-maximal suppression and the technique for gradie?nt computation)
	// are derived from an implementation posted in the following forum (with
	// the
	// clear intent of others using the code):
	// http://forum.java.sun.com/thread.jspa?threadID=546211&start=45&tstart=0
	// My code effectively mimics the algorithm exhibited above.
	// Since I don't know the providence of the code that was posted it is a
	// possibility (though I think a very remote one) that this code violates
	// someone's intellectual property rights. If this concerns you feel free to
	// contact me for an alternative, though less efficient, implementation.
	private void computeGradients(float kernelRadius, int kernelWidth) {
		int initX = kwidth - 1;
		int maxX = width - (kwidth - 1);
		int initY = width * (kwidth - 1);
		int maxY = width * (height - (kwidth - 1));

		int[] imageArray = ImageTools.imageToArray(image);

		// perform convolution in x and y directions
		for (int x = initX; x < maxX; x++) {
			for (int y = initY; y < maxY; y += width) {
				int index = x + y;
				float sumX = imageArray[index] * kernel[0];
				float sumY = sumX;
				int xOffset = 1;
				int yOffset = width;
				for (; xOffset < kwidth;) {
					sumY += kernel[xOffset]
							* (imageArray[index - yOffset] + imageArray[index
									+ yOffset]);
					sumX += kernel[xOffset]
							* (imageArray[index - xOffset] + imageArray[index
									+ xOffset]);
					yOffset += width;
					xOffset++;
				}

				yConv[index] = sumY;
				xConv[index] = sumX;
			}

		}

		for (int x = initX; x < maxX; x++) {
			for (int y = initY; y < maxY; y += width) {
				float sum = 0f;
				int index = x + y;
				for (int i = 1; i < kwidth; i++)
					sum += diffKernel[i]
							* (yConv[index - i] - yConv[index + i]);

				xGradient[index] = sum;
			}

		}

		for (int x = kwidth; x < width - kwidth; x++) {
			for (int y = initY; y < maxY; y += width) {
				float sum = 0.0f;
				int index = x + y;
				int yOffset = width;
				for (int i = 1; i < kwidth; i++) {
					sum += diffKernel[i]
							* (xConv[index - yOffset] - xConv[index + yOffset]);
					yOffset += width;
				}

				yGradient[index] = sum;
			}
		}
	}

	private float gaussian(float x, float sigma) {
		return (float) Math.exp(-(x * x) / (2f * sigma * sigma));
	}

	public float[] getXGradient() {
		return xGradient;
	}

	public float[] getYGradient() {
		return yGradient;
	}

	public BufferedImage getGradientImage() {
		byte[] gradient = new byte[xGradient.length];

		float maxGradient = -1;
		for (int i = 0; i < xGradient.length; i++)
			maxGradient = Math.max(
					maxGradient,
					(float) Math.sqrt(xGradient[i] * xGradient[i]
							+ yGradient[i] * yGradient[i]));

		// normalize the gradient to 0-255
		for (int i = 0; i < xGradient.length; i++) {
			byte g = (byte) (255 * Math.sqrt(xGradient[i] * xGradient[i]
					+ yGradient[i] * yGradient[i]) / maxGradient);
//			if (g > 100)
//				gradient[i] = 0;
//			else if (g > 10)
//				gradient[i] = (byte) 255;

			 gradient[i] = (byte) (g * 10);
		}

		return ImageTools.arrayToImage(gradient, width, height);
	}

	public static void main(String[] args) {
		BufferedImage image = null;
		try {
			image = ImageIO.read(new File("res/raw_16.png"));
		} catch (IOException e) {
			e.printStackTrace();
		}
		Gradient g = new Gradient(image);
		g.process();

		try {
			ImageIO.write(g.getGradientImage(), "png", new File(
					"res/gradient.png"));
		} catch (IOException e) {
			e.printStackTrace();
		}

		System.out.println("Main program finished");
	}
}
